import pandas as pd
import numpy as np
import datetime as dt
import tushare as ts  # 金融接口
from sklearn.svm import SVC
from sklearn.preprocessing import scale  # 数据预处理:标准化
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn import svm, preprocessing

import math
import matplotlib.pyplot as plt


class ParamPredicter():
    # 生成数据
    def initData(self,
                 code="000001.SZ",
                 stockType='daily',
                 start=str(dt.datetime.today() - dt.timedelta(days=200)),
                 end=str(dt.datetime.today() + dt.timedelta(days=1))):
        df = ts.pro_api().query(stockType, ts_code=code, start_date=start, end_date=end)
        stockData = df[['open', 'close', 'high', 'low', 'vol', 'trade_date']].sort_index(ascending=False, ignore_index=True)
        stockData.rename(columns={'trade_date': 'datetime', 'vol': 'volume'}, inplace=True)
        return stockData

    # 传入数据, 参数进行预测
    def predict(self, data, param='close', test_size=0.2):
        origDf = data
        df = origDf[['open', 'close', 'high', 'low', 'volume']]  # 原始数据
        featureData = df[['open', 'high', 'low', 'volume']]  # 特征数据
        # 划分特征值和目标值
        feature = featureData.values
        target = np.array(df[param])

        # 划分训练集，测试集
        feature_train, feature_test, target_train, target_test = train_test_split(feature, target, test_size=test_size)

        pridectedDays = int(math.ceil(test_size * len(df)))  # 预测天数
        lrTool = LinearRegression()
        lrTool.fit(feature_train, target_train)  # 训练
        # 用测试集预测结果
        predictByTest = lrTool.predict(feature_test)  # 测试结果

        # # 组装数据
        index = 0
        # 在前95%的交易日中，设置预测结果和收盘价一致
        while index < len(df) - pridectedDays:
            df.loc[index, param + '_predictedVal'] = origDf.loc[index, param]
            df.loc[index, 'datetime'] = origDf.loc[index, 'datetime']
            index = index + 1

        predictedCnt=0
        # 在后5%的交易日中，用测试集推算预测股价
        while predictedCnt < pridectedDays:
            df.loc[index, param + '_predictedVal'] = predictByTest[predictedCnt]
            df.loc[index, 'datetime'] = origDf.loc[index, 'datetime']
            predictedCnt = predictedCnt + 1
            index = index + 1

        plt.figure()
        # plt.axes(yscale="log")
        df[param + '_predictedVal'].plot(color="red", label='predicted_' + param)
        df[param].plot(color="blue", label='real_' + param)
        plt.legend(loc='best')  # 绘制图例
        major_index = df.index[df.index % 30 == 0]
        major_xtics = df['datetime'][df.index % 30 == 0]
        plt.xticks(major_index, major_xtics)
        plt.setp(plt.gca().get_xticklabels(), rotation=30)
        plt.grid(linestyle='-.')
        plt.show()

    # 传入参数, 进行趋势预测
    def predictTrend(self, data, param='close', test_size=0.2):
        df = data
        # diff列表示本日和上日收盘价的差
        df['diff'] = df["close"] - df["close"].shift(1)
        df['diff'].fillna(0, inplace=True)

        # up列表示本日是否上涨,1表示涨，0表示跌
        # 此处有警告,暂且忽视...
        df['up'] = df['diff']
        df['up'][df['diff'] > 0] = 1
        df['up'][df['diff'] <= 0] = 0

        # 预测值暂且初始化为0
        df['predictForUp'] = 0

        # 目标值是真实的涨跌情况
        target = df['up']
        length = len(df)
        trainNum = int(length * (1-test_size))
        predictNum = length - trainNum
        # 选择指定列作为特征列
        feature = df[['close', 'high', 'low', 'open', 'volume', 'datetime']]
        # 标准化处理特征值
        feature = preprocessing.scale(feature)

        # 训练集的特征值和目标值
        featureTrain = feature[1:trainNum - 1]
        targetTrain = target[1:trainNum - 1]
        svmTool = svm.SVC(kernel='linear')
        svmTool.fit(featureTrain, targetTrain)

        predictedIndex = trainNum
        # 逐行预测测试集
        while predictedIndex < length:
            testFeature = feature[predictedIndex:predictedIndex + 1]
            predictForUp = svmTool.predict(testFeature)
            df.loc[predictedIndex, 'predictForUp'] = predictForUp
            predictedIndex = predictedIndex + 1

        # 该对象只包含预测数据，即只包含测试集
        dfWithPredicted = df[trainNum:length]

        # 开始绘图，创建两个子图
        figure = plt.figure()
        # 创建子图
        (axClose, axUpOrDown) = figure.subplots(2, sharex=True)
        dfWithPredicted['close'].plot(ax=axClose)
        dfWithPredicted['predictForUp'].plot(ax=axUpOrDown, color="red", label='Predicted Data')
        dfWithPredicted['up'].plot(ax=axUpOrDown, color="blue", label='Real Data')
        plt.legend(loc='best')  # 绘制图例
        # 设置x轴坐标标签和旋转角度
        major_index = dfWithPredicted.index[dfWithPredicted.index % 2 == 0]
        major_xtics = dfWithPredicted['datetime'][dfWithPredicted.index % 2 == 0]
        plt.xticks(major_index, major_xtics)
        plt.setp(plt.gca().get_xticklabels(), rotation=30)
        plt.title("预测股票走势")
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.show()


if __name__ == "__main__":
    pd.set_option('mode.chained_assignment', None)
    pd.set_option('display.max_columns', None)  # 展示dataFrame所有数据
    # pd.set_option('display.max_rows', None)  # 展示dataFrame所有数据
    # predict(initData(type='daily'), 'close', 0.1)

    pp = ParamPredicter()
    data = pp.initData(start="20200101", end="20210101")
    pp.predict(data, "open", 0.2)
    pp.predictTrend(data, 0.2)